from cryptography import x509
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.serialization import pkcs12
from cryptography.hazmat.primitives.serialization import BestAvailableEncryption, NoEncryption
import datetime


def generate_private_key():
    """
    生成RSA私钥
    """
    private_key = rsa.generate_private_key(
        public_exponent=65537,
        key_size=2048,
        backend=default_backend()
    )
    return private_key


def generate_csr(private_key, options):
    """
    生成证书签名请求(CSR)
    """
    subject = x509.Name([
        x509.NameAttribute(x509.NameOID.COUNTRY_NAME, options["country_name"]),
        x509.NameAttribute(x509.NameOID.STATE_OR_PROVINCE_NAME, options["state_or_province_name"]),
        x509.NameAttribute(x509.NameOID.LOCALITY_NAME, options["locality_name"]),
        x509.NameAttribute(x509.NameOID.ORGANIZATION_NAME, options["organization_name"]),
        x509.NameAttribute(x509.NameOID.COMMON_NAME, options["common_name"]),
        x509.NameAttribute(x509.NameOID.EMAIL_ADDRESS, options["email_address"]),
    ])
    csr_builder = x509.CertificateSigningRequestBuilder().subject_name(subject)
    csr = csr_builder.sign(private_key, hashes.SHA256(), default_backend())
    return csr


def generate_certificate(private_key, csr, options):
    """
    生成自签名证书
    """
    cert_builder = x509.CertificateBuilder().subject_name(
        x509.Name([
            x509.NameAttribute(x509.NameOID.COMMON_NAME, options["common_name"]),
        ])
    ).issuer_name(csr.subject).public_key(csr.public_key()).serial_number(x509.random_serial_number()
    ).not_valid_before(datetime.datetime.utcnow()).not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=365)
    ).add_extension(x509.BasicConstraints(ca=True, path_length=None), critical=True)
    cert = cert_builder.sign(
        private_key=private_key, algorithm=hashes.SHA256(), backend=default_backend()
    )
    return cert


def save_private_key(private_key, path):
    """
    保存私钥到文件
    """
    print('生成私钥文件...')
    with open(path, "wb") as f:
        f.write(private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=NoEncryption()
        ))


def save_csr(csr, path):
    """
    保存CSR到文件
    """
    print('生成CSR证书文件...')
    with open(path, "wb") as f:
        f.write(csr.public_bytes(serialization.Encoding.PEM))


def save_certificate(cert, path):
    """
    保存PEM格式证书到文件
    """
    print('生成PEM证书文件...')
    with open(path, "wb") as f:
        f.write(cert.public_bytes(serialization.Encoding.PEM))


def save_p12(private_key, cert, options):
    """
    保存P12格式证书到文件
    """
    print('生成P12证书文件...')
    p12 = pkcs12.serialize_key_and_certificates(
        name=options["friendly_name"].encode(),
        encryption_algorithm=BestAvailableEncryption(options["passphrase"].encode()),
        key=private_key,
        cert=cert,
        cas=[cert],
    )
    with open(options["p12_path"], "wb") as f:
        f.write(p12)


def load_private_key(path):
    """
    从文件加载私钥
    """
    with open(path, "rb") as f:
        private_key = serialization.load_pem_private_key(
            f.read(), password=None, backend=default_backend()
        )
    return private_key


def generate(options, save_all=False):
    """
    生成证书
    """
    # if options.get("private_key_path"):
    #     private_key = load_private_key(options["private_key_path"])
    # else:
    #     private_key = generate_private_key()
    private_key = generate_private_key()
    csr = generate_csr(private_key, options)
    cert = generate_certificate(private_key, csr, options)

    save_private_key(private_key, options["private_key_path"])
    save_csr(csr, options["csr_path"])
    save_certificate(cert, options["pem_path"])
    if save_all:
        save_p12(private_key, cert, options)
    print('证书文件已全部生成!')


# 废弃
if __name__ == '__main__':
    # 示例用法
    path = './cert/'
    options = {
        "private_key_path": path + "private.key",
        "csr_path": path + "csr.csr",
        "pem_path": path + "certificate.pem",
        "p12_path": path + "certificate.p12",
        # 通用名称
        "common_name": "yuanxinwenhua",
        # 国家代码
        "country_name": "CN",
        # 省
        "state_or_province_name": "shandong",
        # 地区
        "locality_name": "jinan",
        # 组织名称
        "organization_name": "yuanxinwenhua",
        # 邮箱
        "email_address": "1032428119@qq.com",
        # 证书名称
        "friendly_name": "yuanxinwenhua",
        # 证书密码
        "passphrase": "yuanxinwenhua",
    }

    # 生成所有格式的证书
    # generate(options, save_all=True)

    # 或者单独生成某一个格式的证书
    # generate(options)  # 只生成PEM格式证书
    # generate(options, save_all=True)  # 生成PEM和P12格式证书
